import os
from torch.utils.data import Dataset
from PIL import Image

from typing import List

class ImageListDataset(Dataset):
    r""" a dataset implement that read image filenames from a list file.
    """
    def __init__(self, list_files:List[str], img_dir:str, transform):
        self.transform = transform
        self.img_dir = img_dir
        self.samples = []
        for class_idx, filename in enumerate(list_files):
            with open(os.path.join(img_dir, filename)) as f:
                self.samples.extend((line.strip(), class_idx) for line in f)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path, label = self.samples[idx]

        img = Image.open(os.path.join(self.img_dir, path))
        img = self.transform(img)

        return img, label

